from zipline.api import order_target_percent, symbol, record
import numpy as np

def initialize(context):
    """
    初始化均值回归策略参数
    """
    # 策略参数
    context.period = 20  # 移动平均计算周期
    context.threshold = 1.5  # 触发交易的波动阈值
    context.stop_loss = 0.05  # 5%止损比例
    context.asset = symbol('600519')  # 使用A股标的：贵州茅台
    
    # 风险控制参数
    context.max_position_size = 0.03  # 单笔交易风险控制在3%以内

def handle_data(context, data):
    """
    处理每日数据，实现均值回归策略核心逻辑
    """
    # 获取历史价格数据
    prices = data.history(
        context.asset, 
        'price', 
        bar_count=context.period, 
        frequency='1d'
    )
    
    # 计算均值和标准差
    ma = np.mean(prices)
    std = np.std(prices)
    current_price = data.current(context.asset, 'price')
    
    # 计算当前持仓
    position = context.portfolio.positions[context.asset].amount
    
    # 交易信号逻辑
    if current_price < ma - context.threshold * std:
        # 价格低于下轨，买入信号
        if position == 0:  # 无持仓时开仓
            order_target_percent(context.asset, context.max_position_size)
    elif current_price > ma + context.threshold * std:
        # 价格高于上轨，卖出信号
        if position == 0:  # 无持仓时开仓
            order_target_percent(context.asset, -context.max_position_size)
    
    # 止损逻辑
    if position != 0:
        cost_basis = context.portfolio.positions[context.asset].cost_basis
        if abs(current_price - cost_basis) / cost_basis >= context.stop_loss:
            # 达到止损线，平仓
            order_target_percent(context.asset, 0)
    
    # 每日收盘前平仓不符合持仓条件的头寸
    if position != 0 and not (current_price < ma - context.threshold * std or 
                             current_price > ma + context.threshold * std):
        order_target_percent(context.asset, 0)
    
    # 记录数据（可选）
    record(
        price=current_price,
        moving_avg=ma,
        lower_band=ma - context.threshold * std,
        upper_band=ma + context.threshold * std
    )